import numpy as np
import time
import platform
import os
from qp_taichi import QPSolver

def check_gpu_available():
    """检查是否有GPU可用"""
    # 检查CUDA环境变量
    cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
    if cuda_visible == "-1":
        print("CUDA_VISIBLE_DEVICES设为-1，GPU已禁用")
        return False
        
    try:
        # 尝试导入torch来检测CUDA
        import torch
        if torch.cuda.is_available():
            print(f"检测到{torch.cuda.device_count()}个CUDA设备")
            for i in range(torch.cuda.device_count()):
                print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
            return True
    except ImportError:
        # 如果没有torch，尝试其他方法
        pass
        
    # 检查系统环境
    system = platform.system()
    if system == "Windows":
        try:
            import subprocess
            output = subprocess.check_output("nvidia-smi", shell=True)
            return "NVIDIA" in str(output)
        except:
            pass
    elif system == "Linux":
        try:
            import subprocess
            output = subprocess.check_output("nvidia-smi", shell=True)
            return "NVIDIA" in str(output)
        except:
            pass
            
    print("未检测到可用的GPU，将使用CPU")
    return False

def main():
    """
    使用QP-Taichi求解一个简单的二次规划问题
    
    问题形式:
        min  0.5 * x^T P x + q^T x
        s.t. Gx <= h
             Ax = b
    """
    print("QP-Taichi示例 - 求解简单二次规划问题")
    
    # 定义问题参数 - 使用float32类型
    P = np.array([[4.0, 1.0], [1.0, 2.0]], dtype=np.float32)  # 二次项矩阵，必须是半正定的
    q = np.array([1.0, 1.0], dtype=np.float32)                # 线性项向量
    
    # 不等式约束 Gx <= h
    G = np.array([[1.0, 1.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float32)
    h = np.array([1.0, 0.5, 0.5], dtype=np.float32)
    
    # 等式约束 Ax = b
    A = np.array([[1.0, 1.0]], dtype=np.float32)
    b = np.array([1.0], dtype=np.float32)
    
    # 检查GPU可用性并创建求解器
    use_gpu = check_gpu_available()
    
    try:
        print("\n创建求解器...")
        solver = QPSolver(arch="gpu" if use_gpu else "cpu")
    except Exception as e:
        print(f"求解器初始化失败: {e}")
        print("尝试使用CPU...")
        solver = QPSolver(arch="cpu")
    
    # 设置求解参数
    verbose = True
    
    # 求解
    print("\n开始求解...")
    t_start = time.time()
    
    try:
        result = solver.solve(
            P, q, G, h, A, b, 
            verbose=verbose,
            max_iter=1000,  # 最大迭代次数
            abs_tol=1e-6,   # 更严格的收敛容差
            rel_tol=1e-5,   # 更严格的相对容差
            rho=0.1         # 调小初始rho值
        )
        
        t_end = time.time()
        
        # 输出结果
        print("\n求解结果:")
        print(f"最优解: x = {result.x}")
        print(f"目标函数值: {result.obj_val:.6f}")
        print(f"迭代次数: {result.iterations}")
        print(f"状态: {result.status}")
        print(f"残差 - 原问题: {result.residuals['primal']:.6e}, 对偶: {result.residuals['dual']:.6e}")
        print(f"求解时间: {result.solve_time:.6f} 秒")
        
        # 验证结果
        print("\n验证结果:")
        x = result.x
        obj_value = 0.5 * x.dot(P).dot(x) + q.dot(x)
        print(f"手动计算的目标函数值: {obj_value:.6f}")
        
        # 验证约束
        print("\n约束检查:")
        print(f"不等式约束 Gx <= h:")
        for i in range(G.shape[0]):
            lhs = G[i].dot(x)
            rhs = h[i]
            status = "满足" if lhs <= rhs + 1e-6 else "不满足"
            print(f"  约束 {i+1}: {lhs:.6f} <= {rhs:.6f} - {status}")
        
        print(f"等式约束 Ax = b:")
        for i in range(A.shape[0]):
            lhs = A[i].dot(x)
            rhs = b[i]
            status = "满足" if abs(lhs - rhs) <= 1e-6 else "不满足"
            print(f"  约束 {i+1}: {lhs:.6f} = {rhs:.6f} - {status}")
    
    except Exception as e:
        print(f"求解过程中出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 